from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

from common.llm_builder import llm_messages, llm_no_stream


# 定义工具方法
@tool()
def add(a: int, b: int) -> int:
    """计算两个数字的和"""
    return a + b


# 1-创建一个ChatPromptTemplate
chat_prompt_template = ChatPromptTemplate.from_messages(llm_messages)
print("-----chat_prompt_template_type:", type(chat_prompt_template))
print("-----llm_type:", type(llm_no_stream))

# 2-调用模型(模板+变量=>提示词)
chat_prompt_detail = chat_prompt_template.format_messages(
    role="数学",
    domain="数字计算",
    question="1+150=?"
)
print("-----chat_prompt_detail:", chat_prompt_detail)

# 3-llm添加可用工具
llm_func_tool_dict = {
    "add": add
}
llm_bind_tool = llm_no_stream.bind_tools(tools=[add])
print("-----llm_bind_tool:", llm_bind_tool)
print("-----llm_bind_tool_type:", type(llm_bind_tool))

# 4-构建链式对象
llm_chain = chat_prompt_template | llm_bind_tool

# 5-调用模型(记进行模板赋值)
chat_prompt_detail = llm_chain.invoke({
    "role": "数学",
    "domain": "数字计算",
    "question": "1+150？"
})
print("-----chat_prompt_detail:", chat_prompt_detail)

# 6-打印响应结果
print("-----Hold on, LLM 正在回答！-----")
print(chat_prompt_detail)

# 7-获取对应的参数
for tool_call in chat_prompt_detail.tool_calls:
    print("-----tool_call:", tool_call)
    # 7.1 打印-工具的名称
    tool_call_name = tool_call["name"]
    print("----------tool_call_name:", tool_call_name)

    # 7.2 打印-工具的参数
    tool_call_args = tool_call["args"]
    for tool_call_arg in tool_call_args:
        print("----------tool_call_arg:", tool_call_arg)

    # 7.3 调用tool并传递对应的参数
    tool_call_func = llm_func_tool_dict[tool_call_name]
    tool_call_result = tool_call_func.invoke(tool_call_args)
    print("----------tool_call_result:", tool_call_result)
